package cn.echoparrot.repository.neo4j;

import cn.echoparrot.domain.entity.Dataset;
import cn.echoparrot.domain.entity.Segment;
import org.neo4j.driver.Value;
import org.springframework.data.neo4j.repository.Neo4jRepository;
import org.springframework.data.neo4j.repository.query.Query;
import org.springframework.stereotype.Repository;

import java.util.List;

@Repository
public interface DatasetMapper extends Neo4jRepository<Dataset, Long> {

    @Query("""
            CREATE VECTOR INDEX $indexName IF NOT EXISTS 
            FOR (n:`:#{literal(#labelName)}`) ON (n.embedding)
            OPTIONS {indexConfig: {
            `vector.dimensions`: $dimensions,
            `vector.similarity_function`: 'cosine'
            }}
            """)
    void createVectorIndex(String indexName, String labelName, int dimensions);

    @Query("""
            DROP INDEX $indexName
            """)
    void deleteVectorIndex(String indexName);

    @Query("""
            CALL db.index.vector.queryNodes($indexName, $topK, $embeddingValue)
            YIELD node, score
            WHERE score >= $threshold
            RETURN id(node) AS id,node.text AS text, node.createdAt AS createdAt, score
            """)
    List<Segment> similaritySearch(String indexName, int topK, Value embeddingValue, double threshold);

    /**
     * 如果没有先删除知识库内的文件而先删除还有文件的知识库，那么就会报错。
     * @param datasetId 知识库的唯一标识
     * @return
     */
    @Query("MATCH (d:Dataset) WHERE id(d) = $datasetId DELETE d RETURN count(d)")
    Integer delete(Long datasetId);

    @Query("""
                    MATCH (d:Dataset) 
                    OPTIONAL MATCH (d)-[r:HAS_FILE]->() 
                    RETURN d, count(r) AS fileCount 
                    ORDER BY d.createdAt DESC
                    """)
    List<Dataset> listDatasets();

    @Query("""
                    MATCH (d:Dataset) 
                    WHERE id(d) = $id
                    SET d.name = $dataset.name, d.description = $dataset.description , d.userId = $dataset.userId
                    RETURN d
                    """)
    Dataset update(Dataset dataset);
}